from typing import Dict
import matplotlib.pyplot as plt
from common.app_registry import AppRegistry as AR
from common.app_manager import AppManager
from mmlm.ce.control_loop import ControlLoop, SequentialChain, IterativeRefiner, ConditionalBrancher, SelfCritique, ExternalValidation

class ControlLoopApp(object):
    def __init__(self):
        self.name = 'apps.ces.control_loop_app.ControlLoopApp'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'上下文控制结构......v0.0.1')
        if not AR.is_initialized():
            AR.initialize()
        if params['task'] == 1:
            ControlLoopApp.example_sequential_chain(params=params)
        elif params['task'] == 2:
            ControlLoopApp.example_iterative_refiner(params=params)
        elif params['task'] == 3:
            ControlLoopApp.example_conditional_brancher(params=params)
        elif params['task'] == 4:
            ControlLoopApp.example_self_critique(params=params)
        elif params['task'] == 5:
            ControlLoopApp.example_external_validation(params=params)

    @staticmethod
    def example_sequential_chain(params:Dict = {}) -> None:
        """Example of a sequential chain for data analysis."""
        steps = [
            {
                "name": "extract_entities",
                "prompt_template": "Extract the main entities (people, places, organizations) from this text. For each entity, provide a brief description.\n\nText: {input}",
                "system_message": "You are an expert at extracting and categorizing named entities from text."
            },
            {
                "name": "analyze_relationships",
                "prompt_template": "Based on these entities, analyze the relationships between them:\n\n{input}",
                "system_message": "You are an expert at analyzing relationships between entities."
            },
            {
                "name": "generate_report",
                "prompt_template": "Create a concise summary report based on this relationship analysis:\n\n{input}",
                "system_message": "You are an expert at creating clear, concise reports."
            }
        ]
        chain = SequentialChain(steps=steps, verbose=True)
        sample_text = """
        In 1995, Jeff Bezos founded Amazon in Seattle. Initially an online bookstore, 
        Amazon expanded rapidly under Bezos' leadership. By 2021, Amazon had become 
        one of the world's most valuable companies, and Bezos had briefly overtaken 
        Elon Musk as the world's richest person. Musk, the CEO of Tesla and SpaceX, 
        later reclaimed the top spot after Tesla's stock surged. Meanwhile, Microsoft, 
        founded by Bill Gates in Albuquerque in 1975, continued to be a major tech 
        competitor under CEO Satya Nadella.
        """
        final_output, all_outputs = chain.run(sample_text)
        # Display results
        chain.display_chain_results(all_outputs)
        # Visualize metrics
        chain.visualize_metrics()
        return final_output, all_outputs 
    
    @staticmethod
    def example_iterative_refiner(params:Dict = {}) -> None:
        """Example of iterative refinement for essay writing."""
        # Define a stopping condition based on a quality threshold
        def quality_threshold(response, metadata):
            # Stop if response is over 500 tokens and latency is acceptable
            response_tokens = metadata.get("response_tokens", 0)
            latency = metadata.get("latency", 0)
            return response_tokens > 500 and latency < 5.0
        refiner = IterativeRefiner(
            max_iterations=3,
            stopping_condition=quality_threshold,
            verbose=True
        )
        prompt = "Write a short essay on the future of artificial intelligence."
        final_response, refinement_history = refiner.run(prompt)
        # Display results
        refiner.display_refinement_history(refinement_history)
        # Visualize metrics
        refiner.visualize_metrics()
        return final_response, refinement_history
    
    @staticmethod
    def example_conditional_brancher(params:Dict = {}) -> None:
        """Example of conditional branching for query routing."""
        branches = {
            "technical": {
                "prompt_template": "Provide a technical, detailed explanation of this topic for an expert audience:\n\n{input}",
                "system_message": "You are a technical expert who provides detailed, precise explanations."
            },
            "simplified": {
                "prompt_template": "Explain this topic in simple terms that a 10-year-old would understand:\n\n{input}",
                "system_message": "You are an educator who explains complex topics in simple, accessible language."
            },
            "practical": {
                "prompt_template": "Provide practical, actionable advice on this topic:\n\n{input}",
                "system_message": "You are a practical advisor who provides concrete, actionable guidance."
            }
        }
        brancher = ConditionalBrancher(branches=branches, verbose=True)
        queries = [
            "How does quantum computing work?",
            "What is climate change?",
            "How can I improve my public speaking skills?"
        ]
        results = []
        for query in queries:
            response, run_details = brancher.run(query)
            results.append((query, response, run_details))
            
            # Display results
            brancher.display_branching_results(run_details)
        
        # Visualize metrics
        brancher.visualize_metrics()
        return results
    
    @staticmethod
    def example_self_critique(params:Dict = {}) -> None:
        """Example of self-critique for fact-checking."""
        critique = SelfCritique(
            critique_template="""
            Answer the following question with factual information:

            Question: {input}
            
            Step 1: Write an initial response with all the information you think is relevant.
            
            Step 2: Critically review your response. Check for:
            - Factual errors or inaccuracies
            - Missing important information
            - Potential biases or one-sided perspectives
            - Areas where you're uncertain and should express less confidence
            
            Step 3: Write an improved final response that addresses the issues identified in your critique.
            """,
            verbose=True
        )
        query = "What were the major causes of World War I and how did they lead to the conflict?"
        final_response, run_details = critique.run(query)
        # Display results
        critique.display_results(run_details)
        # Visualize metrics
        critique.visualize_metrics()
        return final_response, run_details
    
    @staticmethod
    def example_external_validation(params:Dict = {}) -> None:
        """Example of external validation for code generation."""
        # Simple validator function that checks for Python syntax errors
        def python_validator(code_response):
            # Extract code blocks
            import re
            code_blocks = re.findall(r"```python(.*?)```", code_response, re.DOTALL)
            if not code_blocks:
                return False, "No Python code blocks found in the response."
            # Check each block for syntax errors
            for i, block in enumerate(code_blocks):
                try:
                    compile(block, "<string>", "exec")
                except SyntaxError as e:
                    return False, f"Syntax error in code block {i+1}: {str(e)}"
            return True, "Code syntax is valid."
        validator = ExternalValidation(
            validator_fn=python_validator,
            max_attempts=3,
            verbose=True
        )
        prompt = "Write a Python function to check if a string is a palindrome."
        final_response, run_details = validator.run(prompt)
        # Display results
        validator.display_results(run_details)
        # Visualize metrics
        validator.visualize_metrics()
        return final_response, run_details